package math;

import java.util.*;
import java.util.regex.*;
import java.util.function.*;
import java.math.*;
import java.util.logging.*;

public class Calculator
{
    public static void main(String[] args)
    {
        Logger EFNLogger = Logger.getLogger("EscapedFunctionNode");
        EFNLogger.setLevel(Level.OFF);

        String expression = "(5 - 3) * 3 ^ 2 + 8 / 4 * 2";
        Manager manager = new Manager(new CalculatorAssistant.Parser(new CalculatorAssistant.Lexer(expression)));
        try(var in = new Scanner(System.in))
        {
            if(args.length > 0)
            {
                expression = args[0];
                MathNode ast = (MathNode)manager.parse();
                System.out.println(ast.getExpressionString() + " = " + manager.evaluate(ast));
            }
            else
            {
                System.out.println("REPL calculator");
                System.out.println("Usage:");
                System.out.println("Enter a expession such as " + expression);
                System.out.println("Output: ");
                MathNode ast = (MathNode)manager.parse();
                System.out.println(ast.getExpressionString() + " = " + manager.evaluate(ast));
                System.out.println("Enter a command to execute a statement,for examples:");
                System.out.println("@var name = 10 --to define a variable");
                System.out.println("@func f(x) = x ^ 2 --to define a function");
                System.out.println("@set name = 1  or  name = 1 --to assign a variable");
                System.out.println("@EOF --to exit");
                System.out.println("Other supported characteristics:");
                System.out.println("LaTeX like syntax to evaluate some math function");
                System.out.println("such as \\sin{\\pi / 2} = 1.0");
                System.out.println("\\log[10]{10} = 1.0 (the expression in the bracket is optional)");
                System.out.println("\\ln{\\e} = 1.0");
                System.out.println("\\sum{i=1}{10}{i^2} = 385.0");

                 // REPL 循环
                while(true)
                {
                    try
                    {
                        System.out.print("> ");
                        String input = in.nextLine().trim();

                        if(input.equals("@EOF") || input.equals("exit") || input.equals("quit"))
                        {
                            System.out.println("Goodbye!");
                            break;
                        }

                        if(input.isEmpty())
                        {
                            continue;
                        }

                        // 创建新的解析器处理当前输入
                        CalculatorAssistant.Lexer lexer = new CalculatorAssistant.Lexer(input);
                        CalculatorAssistant.Parser parser = new CalculatorAssistant.Parser(lexer);
                        Manager currentManager = new Manager(parser, manager.globalScope);

                        SyntaxNode result = currentManager.parse();

                        if(result instanceof MathNode)
                        {
                            double value = currentManager.evaluate((MathNode)result);
                            System.out.println(((MathNode)result).getExpressionString() + " = " + value);
                        }
                        else if(result instanceof Statement)
                        {
                            // Statement 已经在 parse 时执行了
                            System.out.println("OK");
                        }
                    }
                    catch(Exception e)
                    {
                        System.out.println("Error: " + e.getMessage());
                        // 打印更详细的错误信息用于调试
                        // e.printStackTrace();
                    }
                }
            }
        }
        catch(Exception e)
        {
            System.out.println("Fatal error: " + e.getMessage());
            e.printStackTrace();
        }
    }
}

class CalculatorAssistant
{
    static class Lexer 
    {
        private String input;
        private int index;

        int getIndex()
        {
            return index;
        }

        void backTo(int i)
        {
            if(i <= input.length())
            {
                index = i;
            }
            else
            {
                throw new IllegalStateException("The index backing to is out of bound: beckTo(" + i + ")");
            }
        }

        static int EOF = -1;
        static int INVALID = 0;
        static int NUMBER = 1;
        static int PLUS = 2;
        static int MINUS = 3;
        static int MULTIPLY = 4;
        static int DIVIDE = 5;
        static int POWER = 6;
        static int LPAREN = 7;
        static int RPAREN = 8;
        static int FACTORIAL = 9;
        static int WORD = 10;
        static int LBRACK = 11;
        static int RBRACK = 12;
        static int LBRACE = 13;
        static int RBRACE = 14;
        static int COMMA = 15;
        static int COMMAND = 16;
        static int ASSIGN = 17;
        static int ESCAPED_SEQUENCE = 18;

        public Lexer(String input)
        {
            this.input = input;
            index = 0;
        }

        public Token nextToken()
        {
            skipWhiteSpace();

            if(index >= input.length())
            {
                return new Token("EOF",EOF);
            }

            char currentChar = input.charAt(index);
            if(Character.isDigit(currentChar))
            {
                return number();
            }
            if(Character.isUnicodeIdentifierStart(currentChar))
            {
                return word();
            }
            if(currentChar == '\\')
            {
                consume();// '\\'
                return escaped();
            }
            if(currentChar == '@')
            {
                return command();
            }
            try
            {
                switch(currentChar)
                {
                    case'+':return new Token("+",PLUS);
                    case'-':return new Token("-",MINUS);
                    case'*':return new Token("*",MULTIPLY);
                    case'/':return new Token("/",DIVIDE);
                    case'(':return new Token("(",LPAREN);
                    case')':return new Token(")",RPAREN);
                    case'^':return new Token("^",POWER);
                    case'!':return countFactorialQuantity();
                    case'[':return new Token("[",LBRACK);
                    case']':return new Token("]",RBRACK);
                    case'{':return new Token("{",LBRACE);
                    case'}':return new Token("}",RBRACE);
                    case',':return new Token(",",COMMA);
                    case'=':return new Token("=",ASSIGN);
                }

                return new Token(currentChar + "",INVALID);
            }
            finally
            {
                consume();
            }
        }

        private Token countFactorialQuantity()
        {
            StringBuilder f = new StringBuilder();
            while(index < input.length() && input.charAt(index) == '!')
            {
                consume();
                f.append("!");
            }
            deconsume();//Consumed finally in the method nextToken
            
            return new Token(f.toString(),FACTORIAL);
        }

        private void skipWhiteSpace()
        {
            while(index < input.length() && Character.isWhitespace(input.charAt(index)))
            {
                consume();
            }
        }

        private void deconsume()
        {
            if(index > 0)
            {
                index--;
            }
        }

        private void consume()
        {
            if(index < input.length())
            {
                index++;
            }
        }

        private Token number()
        {
            StringBuilder builder = new StringBuilder();
            boolean isHead = true;
            boolean dotted = false;
            while(index < input.length() && (Character.isDigit(input.charAt(index)) || Character.compare(input.charAt(index),'.') == 0))
            {
                if(!isHead && Character.compare(input.charAt(index),'.') == 0 && !dotted)
                {
                    dotted = true;
                }
                isHead = false;
                builder.append(input.charAt(index));
                consume();
            }

            if(index < input.length() && Character.compare(input.charAt(index),'.') == 0 && dotted == true)
            {
                throw new IllegalArgumentException("Invalid number: " + builder.toString() + ".");
            }

            return new Token(builder.toString(),NUMBER);
        }

        private String wordHelper()
        {
            StringBuilder builder = new StringBuilder();
            if(index < input.length() && Character.isUnicodeIdentifierStart(input.charAt(index)))
            {
                builder.append(input.charAt(index));
                consume();
            }

            while(index < input.length() && Character.isUnicodeIdentifierPart(input.charAt(index)))
            {
                builder.append(input.charAt(index));
                consume();
            }
            return builder.toString();
        }

        private Token word()
        {
            return new Token(wordHelper(),WORD);
        }

        private Token specialSequence(int type,String message)
        {
            consume();

            if(index >= input.length())
            {
                throw new IllegalArgumentException(message);
            }

            String result = wordHelper();

            if(result.length() == 0)
            {
                throw new IllegalArgumentException(message);
            }

            return new Token(result,type);
        }
            
        private Token escaped()
        {
            return specialSequence(ESCAPED_SEQUENCE,"Incompleted escaped sequence found");
        }

        private Token command()
        {
            return specialSequence(COMMAND,"Incompleted command found");
        }
    }

    static class Token
    {
        String value;
        int type;

        public Token(String value,int type)
        {
            this.value = value;
            this.type = type;
        }

        @Override
        public String toString()
        {
            return getClass().getName() + "<value=" + value + ">";
        }

    }

    /**Grammer:
     * parse: statement;
     * statement: expression 
     *          | command 
     *          | assign
     *          ;
     * assign: variable '=' expression;
     *
     * EXPRESSION:
     *
     *   expression: plusLevel ('+' | '-' plusLevel)*;
     *   plusLevel: multiplyLevel ('*' | '/' multiplyLevel)*;
     *   multiplyLevel: powerLevel ('^' multiplyLevel)?;
     *   powerLevel: ('+' | '-' powerLevel)
     *             | (powerLevel '!' | '!!')
     *             | '(' expression ')'
     *             | escaped argumentGroup
     *             | functionCall
     *             | variable
     *             | number
     *             ;
     *   escaped: [a-zA-Z_]+; //Symbol table need
     *   argumentGroup: (optionalArgument | requiredArgument)*;
     *   optionalArgument: '[' expression ']';
     *   requiredArgument: '{' expression '}';
     *                   | functionArgument;
     *                   ;
     *   functionArgument: '(' expression (',' expression)* ')'
     *   functionCall: identifier functionArgument;
     *   variable: identifier;
     *   identifier: [a-zA-Z][a-zA-Z_0-9]*;
     *   number: [0-9]+ ('.' [0-9]+)?;
     *
     * COMMAND:
     *
     *   command: def
     *          | assign
     *          | clear
     *          ;
     *   def: funcDef
     *      | varDef
     *      ;
     *   functionDef: '@func' identifier parameters '=' functionContent;
     *   varDef: '@var' identifier '=' expression;
     *   functionContent: expression;
     *   set: '@set' assign;
     */
    static class Parser
    {
        Lexer lexer;
        Token currentToken;

        //private final Scope globalScope;

        static final Map<Integer,String> OPERATOR_MAP = new HashMap<>(){{
            put(Integer.valueOf(Lexer.PLUS),"+");
            put(Integer.valueOf(Lexer.MINUS),"-");
            put(Integer.valueOf(Lexer.MULTIPLY),"*");
            put(Integer.valueOf(Lexer.DIVIDE),"/");
            put(Integer.valueOf(Lexer.POWER),"^");
            put(Integer.valueOf(Lexer.LPAREN),"(");
            put(Integer.valueOf(Lexer.RPAREN),")");
            put(Integer.valueOf(Lexer.FACTORIAL),"!");
        }};

        static final List<String> SPECIAL_LATEX_FUNCTION = List.of("sum","prod","int");
        static final List<String> LATEX_FUNCTION = List.of("sin","cos","tan","sec","csc","cot","arcsin","arccos","arctan","frac","sqrt","log","ln");
        static final List<String> LATEX_CONSTANT = List.of("e","pi","infty","i");

        public Parser(Lexer lexer)
        {
            this.lexer = lexer;
            currentToken = lexer.nextToken();
        }
        
        public Parser setLexer(Lexer l)
        {
            this.lexer = l;
            currentToken = lexer.nextToken();
            return this;
        }

        private Token lookAhead()
        {
            return currentToken;
        }

        private void match(int tokenType)
        {
            if(currentToken.type == tokenType)
            {
                currentToken = lexer.nextToken();
            }
            else
            {
                throw new IllegalArgumentException("Expected token type: " + tokenType + " ,found: " + currentToken.type);
            }
        }

        //public Parser executeStatement(Statement stat)
        //{
            //return executeStatement(stat,globalScope);
        //}
//
        //public Parser executeStatement(Statement stat,Scope scope)
        //{
            //stat.execute(scope);
        //}

        public SyntaxNode parse()
        {
            checkValid();

            int saved = save();
            Token savedToken = currentToken;
            try
            {
                return statement();
            }
            catch(IllegalArgumentException e)
            {
                lexer.backTo(saved);
                currentToken = savedToken;
                return expression();
            }
        }

        public SyntaxNode statement()
        {
            checkValid();

            if(lookAhead().type == Lexer.COMMAND)
            {
                return command();
            }
            return assign();
        }

        private int save()
        {
            return lexer.getIndex();
        }

        public Assignment assign()
        {
            checkValid();

            var v = identifier();
            if(!(v instanceof VariableReferenceNode))
            {
                throw new IllegalArgumentException(v.getClass().getName() + " can't be reassigned");
            }

            VariableReferenceNode vrn = (VariableReferenceNode)v;
            match(Lexer.ASSIGN);
            MathNode expr = expression();

            return new Assignment(vrn,expr);
        }

        public Statement command()
        {
            checkValid();

            Token comm = lookAhead();
            match(Lexer.COMMAND);
            switch(comm.value)
            {
                case"var":return variableDef();
                case"func":return functionDef();
                case"set":return assign();
                //case"clear":return clearStatement();
            }
            throw new IllegalArgumentException("Unknown command found: " + comm.value);
        }

        public Statement variableDef()
        {
            checkValid();

            Token id = lookAhead();
            match(Lexer.WORD);
            match(Lexer.ASSIGN);
            MathNode expr = expression();

            return new VariableDefinition(id.value,expr);
        }

        public Statement functionDef()
        {
            checkValid();

            Token id = lookAhead();
            match(Lexer.WORD);
            match(Lexer.LPAREN);

            List<String> parameters = new ArrayList<>();
            boolean optionalMode = false;
            List<MathNode> defaultValues = new ArrayList<>();
            boolean first = true;
            while(lookAhead().type == Lexer.WORD)
            {
                if(first)
                {
                    first = false;
                }
                else
                {
                    match(Lexer.COMMA);
                }

                parameters.add(lookAhead().value);
                match(Lexer.WORD);
                try
                {
                    if((lookAhead().type == Lexer.ASSIGN && optionalMode == false) || optionalMode == true)
                    {
                        optionalMode = true;
                        match(Lexer.ASSIGN);
                        MathNode defaultValue = expression();
                        defaultValues.add(defaultValue);
                    }
                    else if(lookAhead().type != Lexer.ASSIGN)
                    {
                        continue;
                    }
                }
                catch(IllegalArgumentException e)
                {
                    throw new IllegalArgumentException("Default parameters must be at the last side of the function definition ");
                }
            }

            match(Lexer.RPAREN);
            match(Lexer.ASSIGN);
            MathNode functionBody = expression();

            return new FunctionDefinition(id.value,parameters,defaultValues,functionBody);
        }

        private void checkValid() 
        {
            if(lookAhead().type == Lexer.INVALID)
            {
                throw new IllegalArgumentException("Invalid token found: " + lookAhead().toString());
            }
        }

        private MathNode binaryHelper(Supplier<MathNode> supplier,List<Integer> categories)
        {
            checkValid();
            MathNode left = supplier.get();
            while(categories.contains(lookAhead().type))
            {
                Token current = lookAhead();
                match(lookAhead().type);
                MathNode right = supplier.get();
                left = new BinaryOperatorNode(left,OPERATOR_MAP.get(current.type),right);
            }
            return left;
        }

        public MathNode expression()
        {
            return binaryHelper(() -> plus(),List.of(Integer.valueOf(Lexer.PLUS),Integer.valueOf(Lexer.MINUS)));
        }

        public MathNode plus()
        {
            return binaryHelper(() -> multiply(),List.of(Integer.valueOf(Lexer.MULTIPLY),Integer.valueOf(Lexer.DIVIDE)));
        }

        public MathNode multiply()
        {
            checkValid();
            MathNode left = power();
            while(lookAhead().type == Lexer.POWER)
            {
                match(Lexer.POWER);
                MathNode right = multiply();
                left = new BinaryOperatorNode(left,OPERATOR_MAP.get(Lexer.POWER),right);
            }
            return left;
        }

        private MathNode powerHelper(Integer matchPreviously,Supplier<MathNode> nodeConstructor,Integer matchFinally)
        {
            if(matchPreviously != null) match(matchPreviously.intValue());
            MathNode result = Objects.requireNonNull(nodeConstructor).get();
            if(matchFinally != null) match(matchFinally.intValue());

            if(lookAhead().type == Lexer.FACTORIAL)
            {
                try
                {
                    result = new UnaryOperatorNode(lookAhead().value,result);
                }
                finally
                {
                    match(lookAhead().type);
                }
            }

            return result;
        }

        public MathNode power()
        {
            checkValid();
            boolean succeeded = false;

            Integer matchPreviously = null;
            Supplier<MathNode> nodeConstructor = null;
            Integer matchFinally = null;
            if(lookAhead().type == Lexer.PLUS || lookAhead().type == Lexer.MINUS)
            {
                succeeded = true;
                Token current = lookAhead();
                matchPreviously = current.type;
                matchFinally = null;
                if(current.type == Lexer.MINUS)
                {
                    nodeConstructor = () -> new UnaryOperatorNode("-",power());
                }
                else
                {
                    nodeConstructor = () -> power();
                }
            }
            else if(lookAhead().type == Lexer.LPAREN)
            {
                succeeded = true;

                matchPreviously = Lexer.LPAREN;
                nodeConstructor = () -> expression();
                matchFinally = Lexer.RPAREN;
            }
            else if(lookAhead().type == Lexer.NUMBER)
            {
                succeeded = true;

                Token current = lookAhead();
                matchPreviously = Lexer.NUMBER;
                nodeConstructor = () -> new NumberNode(Double.parseDouble(current.value));
                matchFinally = null;
            }
            else if(lookAhead().type == Lexer.WORD)
            {
                return identifier();
            }
            else if(lookAhead().type == Lexer.ESCAPED_SEQUENCE)
            {
                return escaped();
            }

            if(succeeded == true)
            {
                return powerHelper(matchPreviously,nodeConstructor,matchFinally);
            }
            else
            {
                throw new IllegalArgumentException("Invalid token found: " + lookAhead().toString());
            }
        }

        public MathNode escaped()
        {
            Token id = lookAhead();
            match(Lexer.ESCAPED_SEQUENCE);

            //\sum,\prod,\int
            if(SPECIAL_LATEX_FUNCTION.contains(id.value))
            {
                //{i=1}
                match(Lexer.LBRACE);
                Assignment lowerBound = assign();
                match(Lexer.RBRACE);

                //{10}
                match(Lexer.LBRACE);
                MathNode upperBound = expression();
                match(Lexer.RBRACE);

                //f(i) or {f(i)}
                MathNode expr;
                if(lookAhead().type == Lexer.LBRACE)
                {
                    match(Lexer.LBRACE);
                    expr = expression();
                    match(Lexer.RBRACE);
                }
                else
                {
                    expr = expression();
                }
                return new SpecialLatexFunction(id.value,lowerBound,upperBound,expr);
            }
            else if(LATEX_FUNCTION.contains(id.value))
            {
                List<MathNode> optionalArguments = new ArrayList<>();
                List<MathNode> necessaryArguments = new ArrayList<>();

                while(lookAhead().type == Lexer.LBRACK)
                {
                    match(Lexer.LBRACK);
                    optionalArguments.add(expression());
                    match(Lexer.RBRACK);
                }
                while(lookAhead().type == Lexer.LBRACE)
                {
                    match(Lexer.LBRACE);
                    necessaryArguments.add(expression());
                    match(Lexer.RBRACE);
                }

                return new LatexFunction(id.value,optionalArguments,necessaryArguments);
            }
            else if(LATEX_CONSTANT.contains(id.value))
            {
                return new LatexConstant(id.value);
            }
            else
            {
                throw new IllegalArgumentException(id.value + " is not a supported escaped latex sequence");
            }
        }


        public MathNode identifier()
        {
            Token id = lookAhead();
            match(Lexer.WORD);

            if(lookAhead().type != Lexer.LBRACE && lookAhead().type != Lexer.LBRACK && lookAhead().type != Lexer.LPAREN)
            {
                return new VariableReferenceNode(id.value);
            }
            else 
            {
                match(Lexer.LPAREN);
                boolean first = true;
                List<MathNode> arguments = new ArrayList<>();
                while(lookAhead().type != Lexer.RPAREN)
                {
                    if(first == true)
                    {
                        first = false;
                    }
                    else
                    {
                        match(Lexer.COMMA);
                    }

                    arguments.add(expression());
                }
                match(Lexer.RPAREN);
                return new FunctionCallNode(id.value,arguments);
            }
        }
    }
}

